#!/usr/bin/env python

# This script analyses framestats logfiles and saves a log with number statistics about given videos.
#
# Inputs(given as arguments to script call): video1_framestats.txt, video2_framestats.txt etc...
# Output: frame_log.txt
#
# To do:
#	Check that x axis is correctly scaled (It seems so but messy code regarding kilobyte or bytes)
#	Implement verifying calculations (?) check if dist's sum to 1 etc...

import numpy as np
import pylab as pl
import os.path as osp
import sys
from matplotlib.font_manager import FontProperties as FP

def new_frame(filep):
	line = '1'
	while not line.startswith('[/FRAME]'):
		line = filep.readline()
		if line.startswith('pkt_pos='):
			pos = int(line[line.find('=')+1:])
		if line.startswith('pict_type='):
			pict_type = line[line.find('=')+1:-1]
	return pos, pict_type

# Feed this function with a filepointer and a reference to the plot, it then takes care of plotting. Will be called multiple times with the different logfiles.

def plot(filep,logfile,video_nb):

	pos_type = [] #
	line = '1' #

	# Read data from framestats file
	while line:
		line = f.readline()
		if line.startswith('[FRAME]'):
			position, pict_type = new_frame(f)
			pos_type.append((position, pict_type))

	frame_sizes = [] # list of all frame sizes
	i_size_array = [] # list of all i frame sizes
	p_size_array = [] # list of all p frame sizes
	b_size_array = [] # list of all b frame sizes

	# Populate i,p,b and frame lists
	for i in range(len(pos_type)-1):
		size = pos_type[i+1][0] - pos_type[i][0]
		size = size/1000. # list in kB
		frame_sizes.append(size)
		frame_type = str(pos_type[i][1]).lower()
		if frame_type == 'p':
			p_size_array.append(size)
		elif frame_type == 'i':
			i_size_array.append(size)
		else:
			b_size_array.append(size)


	isize = np.sum(i_size_array) # Size of all i frames [kB]
	psize = np.sum(p_size_array) # Size of all p frames [kB]
	bsize = np.sum(b_size_array) # Size of all p frames [kB]
	totalsize = isize+psize+bsize # total video size [kB]
	totalframes = len(i_size_array)+len(p_size_array)+len(b_size_array) # Total nb of frames

	i_size_mb = isize/1000.
	i_size_mb_str = str("%.1f" % i_size_mb)

	p_size_mb = psize/1000.
	p_size_mb_str = str("%.1f" % p_size_mb)

	b_size_mb = bsize/1000.
	b_size_mb_str = str("%.1f" % b_size_mb)

	# Write useful information to logfile
	logfile.write('Video '+str(video_nb+1)+'\n')
	logfile.write('Nb of frames '+str(totalframes)+'\n')
	logfile.write('Filesize '+str(totalsize/1000.)+' MB'+'\n')
	logfile.write('I-frame data percentage '+str((isize/totalsize)*100)+'\n')
	logfile.write('P-frame data percentage '+str((psize/totalsize)*100)+'\n')
	logfile.write('B-frame data percentage '+str((bsize/totalsize)*100)+'\n')
	logfile.write('I-frame average size '+str(isize/len(i_size_array))+'\n')
	logfile.write('I-frame Max size '+str(max(i_size_array))+'\n')
	logfile.write('I-frame Min size '+str(min(i_size_array))+'\n')
	logfile.write('P-frame average size '+str(psize/len(p_size_array))+'\n')
	logfile.write('P-frame Max size '+str(max(p_size_array))+'\n')
	logfile.write('P-frame Min size '+str(min(p_size_array))+'\n')

	# Avoid division by zero if there are no b-frames in video
	if len(b_size_array)>0:
		logfile.write('B-frame average size '+str(bsize/len(b_size_array))+'\n')
	else:
		logfile.write('B-frame average size '+'0'+'\n')

	logfile.write('I-frames ('+str(len(i_size_array))+' frames, '+i_size_mb_str+' MB)'+'\n')
	logfile.write('P-frames ('+str(len(p_size_array))+' frames, '+p_size_mb_str+' MB)'+'\n')
	logfile.write('B-frames ('+str(len(b_size_array))+' frames, '+b_size_mb_str+' MB)'+'\n')
	logfile.write('GOP nb (average) '+str(totalframes/float(len(i_size_array)))+'\n')
	logfile.write('Bitrate ' +str(float(totalsize)/(totalframes/24))+' kBps')
	logfile.write('\n\n')

if __name__ == '__main__':

	# Open logfile for writing
	logfile=open('data/frame_log.txt','w')

	# Run through all of the given framestats.txt logfiles given in arguments
	for arg in range(len(sys.argv)-1):
		input_filename = sys.argv[arg+1]
		f = open(input_filename,'r')
		plot(f,logfile,arg)

	# Close log file
	logfile.close()
	



